# -*- coding: utf-8 -*-
"""
单模式双波长测试 - 基于多模式多波长光场调制系统框架
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
import os
import time
from datetime import datetime

# 导入自定义模块
from label_utils import create_evaluation_regions_mode_wavelength, evaluate_output, evaluate_all_regions, visualize_labels
from config import Config
from data_generator import MultiModeMultiWavelengthDataGenerator
from visualizer import Visualizer
from trainer import Trainer
from model import MultiModeMultiWavelengthModel
from simulator import Simulator

# 记录开始时间
start_time = time.time()

# 设置随机种子，确保结果可重现
torch.manual_seed(42)
np.random.seed(42)

print("=" * 60)
print("单模式双波长测试 - 基于多模式多波长框架")
print("=" * 60)

# ===== 掩码加载器类（内联定义）=====
class SimpleMaskLoader:
    """简化的相位掩码加载器"""
    
    def __init__(self, config):
        self.config = config
    
    def create_fallback_masks(self, num_layers=2):
        """创建备用聚焦掩码"""
        print("⚠ 创建备用聚焦掩码...")
        
        def create_focusing_mask(size, wavelength, focal_length, pixel_size):
            center = size // 2
            y, x = np.ogrid[:size, :size]
            r_squared = ((x - center) * pixel_size) ** 2 + ((y - center) * pixel_size) ** 2
            k = 2 * np.pi / wavelength
            phase = -k * r_squared / (2 * focal_length)
            return np.mod(phase, 2 * np.pi)
        
        masks = []
        focal_lengths = [50e-3, 100e-3]  # 不同层的焦距
        
        for layer_idx in range(num_layers):
            layer_masks = []
            for wl_idx, wavelength in enumerate(self.config.wavelengths):
                focal_length = focal_lengths[layer_idx % len(focal_lengths)]
                mask = create_focusing_mask(
                    self.config.layer_size, wavelength, focal_length, self.config.pixel_size
                )
                layer_masks.append(mask)
            masks.append(layer_masks)
        
        print(f"✓ 创建了 {len(masks)} 层备用掩码")
        return masks
    
    def get_masks_for_simulation(self, trained_masks=None, num_layers=2):
        """获取用于仿真的掩码"""
        if trained_masks is not None:
            print("✓ 使用训练好的相位掩码进行仿真")
            return trained_masks
        else:
            print("⚠ 使用备用掩码进行仿真")
            return self.create_fallback_masks(num_layers)

# 创建单模式双波长测试配置
config = Config(
    # 基本参数 - 修改为单模式双波长
    num_modes=1,                                # 单模式
    wavelengths=np.array([450e-9, 650e-9]),     # 双波长：蓝光和红光
    
    # 空间参数
    field_size=50,                              # 场大小(像素)
    layer_size=200,                             # 层大小(像素)
    focus_radius=5,                             # 焦点半径(像素)
    detectsize=15,                              # 检测区域大小(像素)
    
    # 物理参数
    z_layers=40e-6,                             # 层间距离(m)
    z_prop=300e-6,                              # 传播距离(m)
    z_step=20e-6,                               # 传播步长(m)
    pixel_size=1e-6,                            # 像素大小(m)
    
    # 检测区域偏移 - 为每个波长定义不同的偏移
    offsets=[(0,0), (30,0)],                    # 两个波长的检测区域偏移
    
    # 训练参数 - 减少训练轮数以便快速测试
    learning_rate=0.01,                         # 学习率
    lr_decay=0.99,                              # 学习率衰减
    epochs=100,                                 # 训练轮数（减少用于测试）
    batch_size=1,                               # 批量大小（单模式）
    
    # 保存参数
    save_dir="./test_single_mode_dual_wavelength/",  # 测试保存目录
    flag_savemat=True                           # 是否保存.mat文件
)

print(f"测试配置:")
print(f"  模式数: {config.num_modes}")
print(f"  波长: {[f'{wl*1e9:.0f}nm' for wl in config.wavelengths]}")
print(f"  训练轮次: {config.epochs}")
print(f"  保存目录: {config.save_dir}")

# 确保保存目录存在
os.makedirs(config.save_dir, exist_ok=True)
os.makedirs(os.path.join(config.save_dir, "trained_models"), exist_ok=True)

# ===== 阶段1: 数据准备 =====
print("\n" + "="*50)
print("阶段1: 数据准备")
print("="*50)

# 创建数据生成器
print("创建数据生成器...")
data_generator = MultiModeMultiWavelengthDataGenerator(config)

# 生成单模式双波长标签
print("生成标签...")
labels = data_generator.generate_labels()

# 可视化标签布局
print("可视化标签布局...")
visualize_labels(labels, config.wavelengths)

# 创建评估区域
print("创建评估区域...")
evaluation_regions = create_evaluation_regions_mode_wavelength(
    config.layer_size, 
    config.layer_size, 
    config.focus_radius, 
    detectsize=config.detectsize
)

print("✓ 数据准备完成")

# ===== 阶段2: 模型训练 =====
print("\n" + "="*50)
print("阶段2: 模型训练")
print("="*50)

# 检查是否存在已训练的模型
trained_models_dir = os.path.join(config.save_dir, "trained_models")
existing_models = []
if os.path.exists(trained_models_dir):
    model_files = [f for f in os.listdir(trained_models_dir) if f.startswith("model_") and f.endswith("layers.pth")]
    existing_models = [f for f in model_files]

# 定义要训练的层数选项（简化测试）
num_layer_options = [1, 2]

# 询问是否使用现有模型或重新训练
use_existing = False
if existing_models:
    print(f"发现已存在的训练模型: {existing_models}")
    try:
        response = input("是否使用现有模型？(y/n，默认n): ").lower().strip()
        use_existing = response == 'y'
    except:
        print("使用默认选项：重新训练")
        use_existing = False

if use_existing and existing_models:
    print("加载现有训练模型...")
    results = {'models': [], 'losses': [], 'phase_masks': [], 'weights_pred': [], 'visibility': []}
    
    for num_layers in num_layer_options:
        model_path = os.path.join(trained_models_dir, f"model_{num_layers}layers.pth")
        if os.path.exists(model_path):
            print(f"加载 {num_layers} 层模型...")
            
            try:
                # 加载模型检查点
                checkpoint = torch.load(model_path, map_location='cpu')
                
                # 创建模型实例
                model = MultiModeMultiWavelengthModel(config, num_layers)
                model.load_state_dict(checkpoint['model_state_dict'])
                
                # 提取相位掩码
                phase_masks = []
                if hasattr(model, 'get_phase_masks_for_simulation'):
                    phase_masks = model.get_phase_masks_for_simulation()
                else:
                    # 兼容旧版本
                    for layer in model.layers:
                        phase = layer.phase.detach().cpu().numpy()
                        phase = phase % (2 * np.pi)
                        wavelength_masks = []
                        for _ in range(len(config.wavelengths)):
                            wavelength_masks.append(phase)
                        phase_masks.append(wavelength_masks)
                
                # 获取训练损失和可见度
                losses = checkpoint.get('train_losses', [])
                visibility = checkpoint.get('visibility', [])
                
                # 如果没有可见度数据，需要重新评估
                if not visibility:
                    print(f"  重新评估 {num_layers} 层模型...")
                    trainer_temp = Trainer(config, data_generator, MultiModeMultiWavelengthModel, evaluation_regions=evaluation_regions)
                    test_loader = trainer_temp._create_data_loaders()[1]
                    eval_results = trainer_temp._evaluate_model(model, test_loader)
                    visibility = eval_results['visibility']
                    weights_pred = eval_results['weights_pred']
                else:
                    weights_pred = []
                
                results['models'].append(model)
                results['losses'].append(losses)
                results['phase_masks'].append(phase_masks)
                results['weights_pred'].append(weights_pred)
                results['visibility'].append(visibility)
                
                print(f"✓ 成功加载 {num_layers} 层模型")
                
            except Exception as e:
                print(f"✗ 加载 {num_layers} 层模型失败: {e}")
                # 添加空结果以保持索引一致
                results['models'].append(None)
                results['losses'].append([])
                results['phase_masks'].append([])
                results['weights_pred'].append([])
                results['visibility'].append([])
        else:
            print(f"✗ 未找到 {num_layers} 层模型文件")
            # 添加空结果以保持索引一致
            results['models'].append(None)
            results['losses'].append([])
            results['phase_masks'].append([])
            results['weights_pred'].append([])
            results['visibility'].append([])
else:
    print("开始训练新模型...")
    
    # 创建训练器
    trainer = Trainer(config, data_generator, MultiModeMultiWavelengthModel, evaluation_regions=evaluation_regions)
    
    # 训练多个层数的模型
    results = trainer.train_multiple_models(num_layer_options)

print("✓ 模型准备完成")

# ===== 阶段3: 结果分析 =====
print("\n" + "="*50)
print("阶段3: 结果分析")
print("="*50)

# 检查可见度数据结构
print("检查训练结果...")
print(f"results键: {list(results.keys())}")
print(f"可见度数据结构: {len(results['visibility'])}层")

valid_results = []
for i, vis_data in enumerate(results['visibility']):
    if vis_data:  # 只处理非空的可见度数据
        expected_length = len(config.wavelengths) * config.num_modes
        print(f"第{i+1}层 ({num_layer_options[i]}层模型):")
        print(f"  数据长度: {len(vis_data)}")
        print(f"  期望长度: {expected_length} ({len(config.wavelengths)}波长 × {config.num_modes}模式)")
        
        if len(vis_data) == expected_length:
            print(f"  ✅ 数据长度匹配！")
            valid_results.append(i)
            # 按波长显示（单模式）
            vis_array = np.array(vis_data).reshape(len(config.wavelengths), config.num_modes)
            for wl_idx, wl in enumerate(config.wavelengths):
                wl_nm = wl * 1e9
                print(f"    {wl_nm:.0f}nm: 模式1={vis_array[wl_idx, 0]:.6f}")
        else:
            print(f"  ❌ 数据长度不匹配！")
    else:
        print(f"第{i+1}层 ({num_layer_options[i]}层模型): 无可见度数据")

# 可视化训练损失
if results['losses'] and any(results['losses']):
    print("可视化训练损失...")
    plt.figure(figsize=(10, 6))
    for i, num_layers in enumerate(num_layer_options):
        if results['losses'][i]:  # 确保有损失数据
            plt.plot(results['losses'][i], label=f'{num_layers} 层')
    plt.xlabel('训练轮次')
    plt.ylabel('损失值')
    plt.title('单模式双波长 - 不同层数模型的训练损失')
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
    plt.show()

# 可视化双波长可见度结果
if valid_results:
    print("可视化双波长可见度结果...")
    fig, axes = plt.subplots(1, len(config.wavelengths), figsize=(12, 5))
    if len(config.wavelengths) == 1:
        axes = [axes]

    wavelength_names = [f"{wl*1e9:.0f}nm" for wl in config.wavelengths]
    colors = ['blue', 'red']  # 蓝光和红光

    for wl_idx, (ax, wl_name, color) in enumerate(zip(axes, wavelength_names, colors)):
        # 提取每个波长下不同层数模型的可见度
        layer_visibilities = []
        valid_layer_options = []
        
        for layer_idx in valid_results:
            vis_data = results['visibility'][layer_idx]
            if len(vis_data) == len(config.wavelengths) * config.num_modes:
                vis_array = np.array(vis_data).reshape(len(config.wavelengths), config.num_modes)
                wl_vis = vis_array[wl_idx, 0]  # 单模式
                layer_visibilities.append(wl_vis)
                valid_layer_options.append(num_layer_options[layer_idx])
        
        if layer_visibilities:
            ax.plot(valid_layer_options, layer_visibilities, 
                    'o-', color=color, linewidth=2, markersize=8, label='单模式')
        
        ax.set_xlabel('层数')
        ax.set_ylabel('可见度')
        ax.set_title(f'{wl_name} 可见度')
        ax.legend()
        ax.grid(True)
        ax.set_ylim(0, 1)

    plt.suptitle('单模式双波长可见度对比', fontsize=16)
    plt.tight_layout()
    plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
    plt.show()

# 创建可视化器并生成详细分析
if valid_results:
    print("创建详细可视化分析...")
    try:
        visualizer = Visualizer(config)
        
        # 组织数据
        visibility_by_mode = visualizer.organize_visibility_by_mode(results, config, num_layer_options)
        
        # 绘制图表
        visualizer.plot_visibility_by_mode(visibility_by_mode, num_layer_options, 
                                           save_path=f"{config.save_dir}/visibility_by_mode.png")
        
        visualizer.plot_visibility_comparison_by_mode_wavelength(visibility_by_mode, num_layer_options,
                                                                save_path=f"{config.save_dir}/visibility_matrix.png")
        
        # 打印摘要和保存数据
        visualizer.print_visibility_summary(visibility_by_mode, num_layer_options)
        visualizer.save_visibility_data(visibility_by_mode, num_layer_options, f"{config.save_dir}/visibility_data.csv")
    except Exception as e:
        print(f"可视化分析出错: {e}")

print("✓ 结果分析完成")

# ===== 阶段4: 光场传播仿真（所有模型）=====
print("\n" + "="*50)
print("阶段4: 光场传播仿真（所有模型）")
print("="*50)

# 获取所有有效模型
print("获取所有有效模型...")
valid_models = [(i, results['models'][i], results['visibility'][i], num_layer_options[i]) 
                for i in range(len(results['models'])) 
                if results['models'][i] is not None and results['visibility'][i]]

if valid_models:
    print(f"发现 {len(valid_models)} 个有效模型，将分别进行传播仿真")
    
    # 创建简化的掩码加载器
    print("准备仿真掩码...")
    mask_loader = SimpleMaskLoader(config)
    
    # 生成输入场（所有模型使用相同的输入场）
    print("生成输入场...")
    input_field = data_generator.generate_input_data()
    
    # 为每个有效模型进行仿真
    for model_idx, (idx, model, visibility, num_layers) in enumerate(valid_models):
        print(f"\n--- 处理第 {model_idx+1}/{len(valid_models)} 个模型: {num_layers}层 ---")
        print(f"平均可见度: {np.mean(visibility):.4f}")
        
        # 获取该模型的相位掩码
        phase_masks = results['phase_masks'][idx]
        
        # 获取仿真掩码
        simulation_masks = mask_loader.get_masks_for_simulation(
            trained_masks=phase_masks, 
            num_layers=num_layers
        )
        
        try:
            # 创建模拟器
            print(f"创建 {num_layers} 层模型的模拟器...")
            simulator = Simulator(config, evaluation_regions=evaluation_regions)
            
            # 为单模式生成专用相位掩膜
            print(f"为 {num_layers} 层模型生成专用相位掩膜...")
            mode_specific_masks = simulator.generate_mode_specific_masks(simulation_masks, config.num_modes)
            
            # 模拟光场传播
            print(f"模拟 {num_layers} 层模型的光场传播...")
            simulator.simulate_propagation(
                simulation_masks, 
                input_field, 
                process_all_modes=True,
                mode_specific_masks=mode_specific_masks
            )
            
            # 打印相位掩膜信息
            print(f"保存 {num_layers} 层模型的相位掩膜信息...")
            if hasattr(model, 'print_phase_masks'):
                # 为每个模型创建单独的保存目录
                model_save_dir = os.path.join(config.save_dir, f"{num_layers}layer_model")
                os.makedirs(model_save_dir, exist_ok=True)
                model.print_phase_masks(save_path=model_save_dir)
            
            print(f"✓ {num_layers} 层模型仿真完成")
            
        except Exception as e:
            print(f"❌ {num_layers} 层模型仿真过程出错: {e}")
            print("继续处理下一个模型...")
            continue
    
    # 生成所有模型的性能对比总结
    print("\n--- 所有模型仿真结果总结 ---")
    print("模型性能对比:")
    
    for model_idx, (idx, model, visibility, num_layers) in enumerate(valid_models):
        avg_visibility = np.mean(visibility)
        print(f"  {num_layers}层模型: 平均可见度 = {avg_visibility:.4f}")
        
        # 显示各波长的详细可见度（单模式）
        vis_array = np.array(visibility).reshape(len(config.wavelengths), config.num_modes)
        for wl_idx, wavelength in enumerate(config.wavelengths):
            wl_nm = wavelength * 1e9
            print(f"    {wl_nm:.0f}nm: {vis_array[wl_idx, 0]:.4f}")
    
    # 找出最佳模型
    best_model_info = max(valid_models, key=lambda x: np.mean(x[2]))
    best_idx, best_model, best_visibility, best_num_layers = best_model_info
    
    print(f"\n🏆 最佳模型: {best_num_layers}层")
    print(f"   平均可见度: {np.mean(best_visibility):.4f}")
    
    # 显示最佳模型各波长的可见度
    vis_array = np.array(best_visibility).reshape(len(config.wavelengths), config.num_modes)
    for wl_idx, wavelength in enumerate(config.wavelengths):
        wl_nm = wavelength * 1e9
        print(f"   {wl_nm:.0f}nm: {vis_array[wl_idx, 0]:.4f}")
    
    print(f"\n✅ 所有 {len(valid_models)} 个模型的仿真完成")
    
else:
    print("❌ 没有有效的训练结果，跳过仿真")

print("✓ 光场传播仿真阶段完成")

# ===== 阶段5: 保存最终结果 =====
print("\n" + "="*50)
print("阶段5: 保存最终结果")
print("="*50)

# 保存完整结果
print("保存完整结果...")
timestamp = datetime.now().strftime("%Y%m%d_%H%M")

# 准备保存数据
save_data = {
    'config': config.__dict__,  # 保存配置字典
    'models_state_dict': [model.state_dict() if model else {} for model in results['models']],
    'losses': results['losses'],
    'visibility': results['visibility'],
    'num_layer_options': num_layer_options,
    'timestamp': timestamp,
    'training_completed': True,
    'test_type': 'single_mode_dual_wavelength'
}

# 如果有最佳模型，添加最佳模型信息
if 'best_idx' in locals():
    save_data.update({
        'best_model_idx': best_idx,
        'best_num_layers': best_num_layers,
        'best_avg_visibility': np.mean(best_visibility) if best_visibility else 0
    })

# 保存主结果文件
main_results_path = f"{config.save_dir}/complete_results_{timestamp}.pth"
torch.save(save_data, main_results_path)
print(f"✓ 主结果已保存: {main_results_path}")

# 保存单独的模型文件（便于后续加载）
for i, (model, num_layers) in enumerate(zip(results['models'], num_layer_options)):
    if model is not None:
        model_save_path = os.path.join(config.save_dir, "trained_models", f"model_{num_layers}layers.pth")
        torch.save({
            'model_state_dict': model.state_dict(),
            'model_config': {
                'num_layers': num_layers,
                'model_class': 'MultiModeMultiWavelengthModel'
            },
            'train_losses': results['losses'][i] if i < len(results['losses']) else [],
            'visibility': results['visibility'][i] if i < len(results['visibility']) else [],
            'config': config.__dict__,
            'timestamp': timestamp,
            'test_type': 'single_mode_dual_wavelength'
        }, model_save_path)
        print(f"✓ {num_layers}层模型已保存: {model_save_path}")

print("✓ 所有结果保存完成")

# ===== 程序完成 =====
print("\n" + "="*60)
print("单模式双波长测试完成!")
print("="*60)

# 计算总执行时间
total_end_time = time.time()
total_time = total_end_time - start_time
print(f"总执行时间: {total_time:.2f} 秒 ({total_time/60:.2f} 分钟)")

# 打印最终摘要
print("\n=== 测试执行摘要 ===")
print(f"测试类型: 单模式双波长")
print(f"模式数: {config.num_modes}")
print(f"波长: {[f'{wl*1e9:.0f}nm' for wl in config.wavelengths]}")
print(f"训练层数选项: {num_layer_options}")
print(f"训练轮次: {config.epochs}")
print(f"保存目录: {config.save_dir}")

if 'best_idx' in locals():
    print(f"最佳模型: {best_num_layers}层")
    print(f"最佳平均可见度: {np.mean(best_visibility):.4f}")
    # 显示最佳模型各波长详细结果
    vis_array = np.array(best_visibility).reshape(len(config.wavelengths), config.num_modes)
    for wl_idx, wavelength in enumerate(config.wavelengths):
        wl_nm = wavelength * 1e9
        print(f"  {wl_nm:.0f}nm: {vis_array[wl_idx, 0]:.4f}")

print(f"结果文件: complete_results_{timestamp}.pth")
print("="*40)

============================================================
单模式双波长测试 - 基于多模式多波长框架
============================================================
测试配置:
  模式数: 1
  波长: ['450nm', '650nm']
  训练轮次: 100
  保存目录: ./test_single_mode_dual_wavelength/

==================================================
阶段1: 数据准备
==================================================
创建数据生成器...
生成标签...
可视化标签布局...

/home/shiyue/ODNN/ODNN_WAVE/label_utils.py:179: UserWarning: Glyph 27169 (\N{CJK UNIFIED IDEOGRAPH-6A21}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/home/shiyue/ODNN/ODNN_WAVE/label_utils.py:179: UserWarning: Glyph 24335 (\N{CJK UNIFIED IDEOGRAPH-5F0F}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 27169 (\N{CJK UNIFIED IDEOGRAPH-6A21}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 24335 (\N{CJK UNIFIED IDEOGRAPH-5F0F}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)

创建评估区域...

/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 35780 (\N{CJK UNIFIED IDEOGRAPH-8BC4}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 20272 (\N{CJK UNIFIED IDEOGRAPH-4F30}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21306 (\N{CJK UNIFIED IDEOGRAPH-533A}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 22495 (\N{CJK UNIFIED IDEOGRAPH-57DF}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 27874 (\N{CJK UNIFIED IDEOGRAPH-6CE2}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 38271 (\N{CJK UNIFIED IDEOGRAPH-957F}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 32452 (\N{CJK UNIFIED IDEOGRAPH-7EC4}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21512 (\N{CJK UNIFIED IDEOGRAPH-5408}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)

✓ 数据准备完成

==================================================
阶段2: 模型训练
==================================================
发现已存在的训练模型: ['model_2layers.pth', 'model_1layers.pth']
开始训练新模型...
使用外部提供的评估区域: 9个区域

==================================================
开始训练 1 层模型...
==================================================
原始数据形状: (50, 50, 55)
选择后的数据形状: (1, 50, 50)
模式 1 振幅范围: 0.0 - 0.10288702696561813
计算波长系数 (基准波长: 650.0nm)
  波长 450.0nm: 系数 = 1.4444
  波长 650.0nm: 系数 = 1.0000
开始训练 - 设备: cuda:0
训练参数: epochs=100, lr=0.01
Epoch [0/100], Loss: 0.002080081263557076
✓ 训练完成!
计算可见度: 2波长, 1批次, 9区域, 1模式
处理波长 450nm (索引0)
  处理模式 1
    批次0: 能量=[np.float32(30.543804), np.float32(0.37201512), np.float32(0.57639337), np.float32(0.5750987), np.float32(0.98396456), np.float32(0.3484194), np.float32(0.4617597), np.float32(0.28267238), np.float32(0.43021196)], 可见度=0.981660
  模式1平均可见度: 0.981660
处理波长 650nm (索引1)
  处理模式 1
    批次0: 能量=[np.float32(0.13588485), np.float32(0.060532063), np.float32(0.07746573), np.float32(0.34240484), np.float32(10.184021), np.float32(0.32938528), np.float32(0.2760628), np.float32(0.10256021), np.float32(0.05936178)], 可见度=0.988410
  模式1平均可见度: 0.988410
总共计算了 2 个可见度值
期望值: 2
✓ 完整模型已保存到: ./test_single_mode_dual_wavelength/trained_models/trained_model_1layers.pth
✓ 训练好的相位掩码已保存到: ./test_single_mode_dual_wavelength/trained_models/trained_phase_masks_1layers.npz
✓ 训练损失已保存到: ./test_single_mode_dual_wavelength/trained_models/training_losses_1layers.npy

====== 模型所有相位掩膜信息 ======

== 层 0 相位掩膜信息 ==
基础相位掩膜形状: (200, 200)
基础相位掩膜范围: [0.0000, 6.2653]
波长系数: [1.4444444 1.       ]
已保存层 0 的相位掩膜图像到 ./test_single_mode_dual_wavelength/trained_models/phase_mask_visualization_1layers/phase_mask_layer_0.png
✓ 1层模型训练完成，可见度数量: 2

==================================================
开始训练 2 层模型...
==================================================
计算波长系数 (基准波长: 650.0nm)
  波长 450.0nm: 系数 = 1.4444
  波长 650.0nm: 系数 = 1.0000
计算波长系数 (基准波长: 650.0nm)
  波长 450.0nm: 系数 = 1.4444
  波长 650.0nm: 系数 = 1.0000
开始训练 - 设备: cuda:0
训练参数: epochs=100, lr=0.01
Epoch [0/100], Loss: 0.002034377306699753
✓ 训练完成!
计算可见度: 2波长, 1批次, 9区域, 1模式
处理波长 450nm (索引0)
  处理模式 1
    批次0: 能量=[np.float32(46.767437), np.float32(1.0172424), np.float32(0.3915699), np.float32(0.23892052), np.float32(0.38259783), np.float32(0.30100942), np.float32(0.3479494), np.float32(0.23068659), np.float32(0.1959455)], 可见度=0.991655
  模式1平均可见度: 0.991655
处理波长 650nm (索引1)
  处理模式 1
    批次0: 能量=[np.float32(0.15749925), np.float32(9.198645), np.float32(0.28359902), np.float32(0.16443527), np.float32(2.27946), np.float32(0.3945292), np.float32(0.13601878), np.float32(0.10125196), np.float32(0.17875272)], 可见度=0.978225
  模式1平均可见度: 0.978225
总共计算了 2 个可见度值
期望值: 2
✓ 完整模型已保存到: ./test_single_mode_dual_wavelength/trained_models/trained_model_2layers.pth
✓ 训练好的相位掩码已保存到: ./test_single_mode_dual_wavelength/trained_models/trained_phase_masks_2layers.npz
✓ 训练损失已保存到: ./test_single_mode_dual_wavelength/trained_models/training_losses_2layers.npy

====== 模型所有相位掩膜信息 ======

== 层 0 相位掩膜信息 ==
基础相位掩膜形状: (200, 200)
基础相位掩膜范围: [0.0001, 6.2764]
波长系数: [1.4444444 1.       ]
已保存层 0 的相位掩膜图像到 ./test_single_mode_dual_wavelength/trained_models/phase_mask_visualization_2layers/phase_mask_layer_0.png

== 层 1 相位掩膜信息 ==
基础相位掩膜形状: (200, 200)
基础相位掩膜范围: [0.0002, 6.2823]
波长系数: [1.4444444 1.       ]
已保存层 1 的相位掩膜图像到 ./test_single_mode_dual_wavelength/trained_models/phase_mask_visualization_2layers/phase_mask_layer_1.png
✓ 2层模型训练完成，可见度数量: 2
✓ 模型准备完成

==================================================
阶段3: 结果分析
==================================================
检查训练结果...
results键: ['models', 'losses', 'phase_masks', 'weights_pred', 'visibility']
可见度数据结构: 2层
第1层 (1层模型):
  数据长度: 2
  期望长度: 2 (2波长 × 1模式)
  ✅ 数据长度匹配！
    450nm: 模式1=0.981660
    650nm: 模式1=0.988410
第2层 (2层模型):
  数据长度: 2
  期望长度: 2 (2波长 × 1模式)
  ✅ 数据长度匹配！
    450nm: 模式1=0.991655
    650nm: 模式1=0.978225
可视化训练损失...

/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 25439 (\N{CJK UNIFIED IDEOGRAPH-635F}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 22833 (\N{CJK UNIFIED IDEOGRAPH-5931}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 20540 (\N{CJK UNIFIED IDEOGRAPH-503C}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 21333 (\N{CJK UNIFIED IDEOGRAPH-5355}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 27169 (\N{CJK UNIFIED IDEOGRAPH-6A21}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 24335 (\N{CJK UNIFIED IDEOGRAPH-5F0F}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 21452 (\N{CJK UNIFIED IDEOGRAPH-53CC}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 27874 (\N{CJK UNIFIED IDEOGRAPH-6CE2}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 38271 (\N{CJK UNIFIED IDEOGRAPH-957F}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 19981 (\N{CJK UNIFIED IDEOGRAPH-4E0D}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 21516 (\N{CJK UNIFIED IDEOGRAPH-540C}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 23618 (\N{CJK UNIFIED IDEOGRAPH-5C42}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 25968 (\N{CJK UNIFIED IDEOGRAPH-6570}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 22411 (\N{CJK UNIFIED IDEOGRAPH-578B}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 30340 (\N{CJK UNIFIED IDEOGRAPH-7684}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 35757 (\N{CJK UNIFIED IDEOGRAPH-8BAD}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 32451 (\N{CJK UNIFIED IDEOGRAPH-7EC3}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 36718 (\N{CJK UNIFIED IDEOGRAPH-8F6E}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:297: UserWarning: Glyph 27425 (\N{CJK UNIFIED IDEOGRAPH-6B21}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/training_losses.png", dpi=300)

/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 25439 (\N{CJK UNIFIED IDEOGRAPH-635F}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 22833 (\N{CJK UNIFIED IDEOGRAPH-5931}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 20540 (\N{CJK UNIFIED IDEOGRAPH-503C}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21333 (\N{CJK UNIFIED IDEOGRAPH-5355}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21452 (\N{CJK UNIFIED IDEOGRAPH-53CC}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 19981 (\N{CJK UNIFIED IDEOGRAPH-4E0D}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21516 (\N{CJK UNIFIED IDEOGRAPH-540C}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 23618 (\N{CJK UNIFIED IDEOGRAPH-5C42}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 25968 (\N{CJK UNIFIED IDEOGRAPH-6570}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 22411 (\N{CJK UNIFIED IDEOGRAPH-578B}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 30340 (\N{CJK UNIFIED IDEOGRAPH-7684}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 35757 (\N{CJK UNIFIED IDEOGRAPH-8BAD}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 32451 (\N{CJK UNIFIED IDEOGRAPH-7EC3}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 36718 (\N{CJK UNIFIED IDEOGRAPH-8F6E}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 27425 (\N{CJK UNIFIED IDEOGRAPH-6B21}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)

/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 23618 (\N{CJK UNIFIED IDEOGRAPH-5C42}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 25968 (\N{CJK UNIFIED IDEOGRAPH-6570}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 21487 (\N{CJK UNIFIED IDEOGRAPH-53EF}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 35265 (\N{CJK UNIFIED IDEOGRAPH-89C1}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 24230 (\N{CJK UNIFIED IDEOGRAPH-5EA6}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 21333 (\N{CJK UNIFIED IDEOGRAPH-5355}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 27169 (\N{CJK UNIFIED IDEOGRAPH-6A21}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 24335 (\N{CJK UNIFIED IDEOGRAPH-5F0F}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 21452 (\N{CJK UNIFIED IDEOGRAPH-53CC}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 27874 (\N{CJK UNIFIED IDEOGRAPH-6CE2}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 38271 (\N{CJK UNIFIED IDEOGRAPH-957F}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 23545 (\N{CJK UNIFIED IDEOGRAPH-5BF9}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:335: UserWarning: Glyph 27604 (\N{CJK UNIFIED IDEOGRAPH-6BD4}) missing from font(s) DejaVu Sans.
  plt.tight_layout()
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 21487 (\N{CJK UNIFIED IDEOGRAPH-53EF}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 35265 (\N{CJK UNIFIED IDEOGRAPH-89C1}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 24230 (\N{CJK UNIFIED IDEOGRAPH-5EA6}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 23618 (\N{CJK UNIFIED IDEOGRAPH-5C42}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 25968 (\N{CJK UNIFIED IDEOGRAPH-6570}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 21333 (\N{CJK UNIFIED IDEOGRAPH-5355}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 27169 (\N{CJK UNIFIED IDEOGRAPH-6A21}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 24335 (\N{CJK UNIFIED IDEOGRAPH-5F0F}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)

可视化双波长可见度结果...

/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 21452 (\N{CJK UNIFIED IDEOGRAPH-53CC}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 27874 (\N{CJK UNIFIED IDEOGRAPH-6CE2}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 38271 (\N{CJK UNIFIED IDEOGRAPH-957F}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 23545 (\N{CJK UNIFIED IDEOGRAPH-5BF9}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/tmp/ipykernel_62949/1509928505.py:336: UserWarning: Glyph 27604 (\N{CJK UNIFIED IDEOGRAPH-6BD4}) missing from font(s) DejaVu Sans.
  plt.savefig(f"{config.save_dir}/dual_wavelength_visibility.png", dpi=300)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 21487 (\N{CJK UNIFIED IDEOGRAPH-53EF}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 35265 (\N{CJK UNIFIED IDEOGRAPH-89C1}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 24230 (\N{CJK UNIFIED IDEOGRAPH-5EA6}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 23545 (\N{CJK UNIFIED IDEOGRAPH-5BF9}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)
/home/shiyue/miniconda3/envs/ODNN/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Glyph 27604 (\N{CJK UNIFIED IDEOGRAPH-6BD4}) missing from font(s) DejaVu Sans.
  fig.canvas.print_figure(bytes_io, **kw)